import copy
from typing import Optional, List

import math
import torch
import torch.nn.functional as F
from torch import nn, Tensor
from .utils import _get_clones, _get_activation_fn
from module.row_column_decoupled_attention import MultiheadRCDA
from models.module.modules import FFN
from .utils import pos2posemb2d, pos2posemb1d

class TransformerDecoderLayer(nn.Module):
    def __init__(self, d_model=256, d_ffn=1024,
                 dropout=0., activation="relu", n_heads=8,
                 n_levels=3, attention_type="RCDA"):
        super().__init__()

        self.attention_type = attention_type
        self.attention_type = attention_type
        if attention_type=="RCDA":
            attention_module=MultiheadRCDA
        elif attention_type == "nn.MultiheadAttention":
            attention_module=nn.MultiheadAttention
        else:
            raise ValueError(f'unknown {attention_type} attention_type')

        # cross attention
        self.cross_attn = attention_module(d_model, n_heads, dropout=dropout)
        self.dropout1 = nn.Dropout(dropout)
        self.norm1 = nn.LayerNorm(d_model)

        # self attention
        self.self_attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.norm2 = nn.LayerNorm(d_model)


        # level combination
        if n_levels>1:
            self.level_fc = nn.Linear(d_model * n_levels, d_model)

        # ffn
        self.ffn = FFN(d_model, d_ffn, dropout, activation)

    @staticmethod
    def with_pos_embed(tensor, pos):
        return tensor if pos is None else tensor + pos

    def forward(self, tgt, reference_points, srcs, src_padding_masks=None, adapt_pos2d=None,
                adapt_pos1d=None, posemb_row=None, posemb_col=None, posemb_2d=None):
        tgt_len = tgt.shape[1]

        query_pos = pos2posemb2d(reference_points.squeeze(2))
        query_pos = adapt_pos2d(query_pos)
        # self attention
        q = k = self.with_pos_embed(tgt, query_pos)
        tgt2 = self.self_attn(q.transpose(0, 1), k.transpose(0, 1), tgt.transpose(0, 1))[0].transpose(0, 1)
        tgt = tgt + self.dropout2(tgt2)
        tgt = self.norm2(tgt)

        bz, l, c, h, w = srcs.shape
        srcs = srcs.reshape(bz * l, c, h, w).permute(0, 2, 3, 1)



        if self.attention_type == "RCDA":
            query_pos_x = adapt_pos1d(pos2posemb1d(reference_points[..., 0]))
            query_pos_y = adapt_pos1d(pos2posemb1d(reference_points[..., 1]))
            posemb_row = posemb_row.unsqueeze(1).repeat(1, h, 1, 1)
            posemb_col = posemb_col.unsqueeze(2).repeat(1, 1, w, 1)
            src_row = src_col = srcs
            q2 = (tgt + query_pos_x).repeat(l, 1, 1)
            k2 = (tgt + query_pos_y).repeat(l, 1, 1)
            k_row = src_row + posemb_row
            k_col = src_col + posemb_col

            tgt2 = self.cross_attn(q2, k2, k_row, k_col,
                                   srcs, key_padding_mask=src_padding_masks)[0].transpose(0, 1)
        else:
            q2 = (tgt + query_pos).repeat(l, 1, 1).transpose(0, 1)
            k2 =  (srcs + posemb_2d).reshape(bz * l, h * w, c).transpose(0,1)
            v2 = srcs.reshape(bz * l, h * w, c).transpose(0, 1)
            tgt2 = self.cross_attn(q2, k2, v2, key_padding_mask=src_padding_masks.reshape(bz*l, h*w))[0].transpose(0,1)

        if l > 1:
            tgt2 = self.level_fc(tgt2.reshape(bz, l, tgt_len, c).permute(0, 2, 3, 1).reshape(bz, tgt_len, c * l))

        tgt = tgt + self.dropout1(tgt2)
        tgt = self.norm1(tgt)

        # ffn
        tgt = self.ffn(tgt)

        return tgt

